import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf

def accuarcy(output, target, topk=[1,]):
    maxk = max(topk)
    batch_size = target.shape[0]

    pred = tf.math.top_k(output, maxk).indices#前K个最大值所在的序号
    pred = tf.transpose(pred, [1, 0]) #转置
    target_ = tf.broadcast_to(target, pred.shape)#广播
    correct = tf.equal(pred, target_)

    res = []
    for k in topk:
        correct_k = tf.cast(tf.reshape(correct[:k],[-1]), dtype=tf.float32)
        correct_k = tf.reduce_sum(correct_k)
        acc = float(correct_k / batch_size)
        res.append(acc)

    return res


output = tf.random.normal([10,6])
output = tf.math.softmax(output, axis=1)
target = tf.random.uniform([10],maxval=6, dtype=tf.int32)

print('prob:', output.numpy())
pred = tf.argmax(output, axis=1)
print('pred:', pred.numpy())
print('label:', target.numpy())

list = [1,2,3,4,5,6]
acc = accuarcy(output, target, topk=list)
print('top-1-6 acc:',acc)